from IPython.core.display import Image
from IPython.core.display_functions import display
from langgraph.graph import StateGraph
from langgraph.graph import START, END
from pydantic import BaseModel


class CalcState(BaseModel):
    x: int


def addition(state: CalcState) -> CalcState:
    # 注意 这里接受到的是初始化状态
    print(f"init_state: {state}")
    return CalcState(x=state.x + 1)


def subtraction(state: CalcState) -> CalcState:
    # 注意 这里接受到的上一个节点状态
    print(f"init_state: {state}")
    return CalcState(x=state.x - 2)


if __name__ == '__main__':
    # pip install pyppeteer ipython
    builder = StateGraph(CalcState)
    builder.add_node("addition", addition)
    builder.add_node("subtraction", subtraction)
    builder.add_edge(START, "addition")
    builder.add_edge("addition", "subtraction")
    builder.add_edge("subtraction", END)
    print(builder.edges)
    print(builder.nodes)
    graph = builder.compile()
    # graph_image = graph.get_graph(xray=True).draw_mermaid_png()
    # with open("graph.png", "wb") as f:
    #     f.write(graph_image)
    # print("Graph saved to graph.png")
    initial_state = CalcState(x=10)
    x = graph.invoke(initial_state)
    print(x)
